import os
import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset


# 自定义数据集加载器
class IrisDataloader(Dataset):

    def __init__(self, data_path):
        """
        数据加载器
        输入：
        data_path: 数据集路径
        :param data_path:
        """

        # 数据集路径
        self.data_path = data_path
        assert os.path.exists(self.data_path), "dataset does not exits"

        # 加载数据集
        df = pd.read_csv(self.data_path, names=[0, 1, 2, 3, 4])

        # 标签映射
        d = {"Iris-setosa": 0, "Iris-versicolor": 1, "Iris-virginica": 2}
        df[4] = df[4].map(d)

        # 划分特征和标签
        data = df.iloc[:, 0:4]
        label = df.iloc[:, 4:]

        # 数据预处理
        data = (data - np.mean(data)) / np.std(data)  # Z值化

        # 数据类型转换
        self.data = torch.from_numpy(np.array(data, dtype='float32'))
        self.label = torch.from_numpy(np.array(label, dtype='int64'))

        # 数据集大小
        self.data_num = len(label)
        print("当前数据集的大小： ", self.data_num)

    def __len__(self):
        """
        返回数据集大小
        :return: 数据集大小
        """
        return self.data_num

    def __getitem__(self, index):
        """
        返回数据集
        :param index: 索引
        :return: 数据集
        """

        # 数据类型转换
        self.data = list(self.data)
        self.label = list(self.label)

        # 按索引返回数据集
        return self.data[index], self.label[index]
